package com.tgy.algorithm.base._并查集;


import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

class UFNode<V> {
    V value;
    UFNode(V value) {
        this.value = value;
    }
}

class UnionFind<V> {

    // 每一个V对应的Node的map
    private Map<V, UFNode<V>> nodes =  new HashMap<>();
    // 每一个Node对应的Parent Node
    private Map<UFNode<V>, UFNode<V>> parents = new HashMap<>();
    // 并查集里面每一个集合的size
    Map<UFNode<V>, Integer> nodesSize = new HashMap<>();

    UnionFind(List<V> lists){
        for (V v : lists) {
            UFNode<V> node = new UFNode<>(v);
            nodes.put(v, node);
            // 初始化父元素是自身
            parents.put(node, node);
            // 每一个元素自成一个集合
            nodesSize.put(node, 1);
        }
    }

    public UFNode<V> findParent(UFNode<V> node) {
        List<UFNode<V>> pathNodes = new ArrayList<>();
        while (node != parents.get(node)) {
            pathNodes.add(node);
            node = parents.get(node);
        }

        // 扁平化path里面的每一个node，让其父元素直接指向代表元素
        for (UFNode<V> pathNode : pathNodes) {
            parents.put(pathNode, node);
        }

        return node;
    }

    // 查找是否在同一个集合中
    public boolean isSameSet(V v1, V v2) {
        if (!nodes.containsKey(v1) || !nodes.containsKey(v2)) {
            return false;
        }

        UFNode<V> nodeV1 = findParent(nodes.get(v1));
        UFNode<V> nodeV2 = findParent(nodes.get(v2));
        return nodeV1 == nodeV2;
    }

    // 合并
    public void union(V v1, V v2) {
        if (!nodes.containsKey(v1) || !nodes.containsKey(v2)) {
            return;
        }

        UFNode<V> parentV1 = findParent(nodes.get(v1));
        UFNode<V> parentV2 = findParent(nodes.get(v2));
        if (parentV1 == parentV2) {
            return;
        }

        Integer parentSizeV1 = nodesSize.get(parentV1);
        Integer parentSizeV2 = nodesSize.get(parentV2);
        if (parentSizeV1 >= parentSizeV2) {
            parents.put(parentV2, parentV1);
            nodesSize.put(parentV1,parentSizeV1 + parentSizeV2);
            nodesSize.remove(parentV2);
        }else {
            parents.put(parentV1, parentV2);
            nodesSize.put(parentV2,parentSizeV1 + parentSizeV2);
            nodesSize.remove(parentV1);
        }
    }

    public int getMaps() {
        return nodesSize.size();
    }
}


public class _并查集实现 {
    public static void main(String[] args) {

        ArrayList<String> strings = new ArrayList<>();
        strings.add("A");
        strings.add("B");
        strings.add("C");
        strings.add("D");

        UnionFind unionFind = new UnionFind(strings);

        boolean sameSet = unionFind.isSameSet("A", "B");
        System.out.println(sameSet);
        unionFind.union("A", "B");
        sameSet = unionFind.isSameSet("A", "B");
        System.out.println(sameSet);
        unionFind.union("A", "C");
        sameSet = unionFind.isSameSet("A", "C");
        System.out.println(sameSet);
    }
}
